import os
import time


class Config:
    current_time = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
    # save_file = current_time
    is_balanced_weight = False
    img_size = [224, 224]
    initial_learning_rate = 0.0005      # 0.0005
    # warmup_learning_rate = 0.02
    momentum = 0.9
    weight_decay = 1e-4
    train_batch_size = 32               # 32
    test_batch_size = 10
    train_number_epochs = 200
    num_workers = 8
    network = "resnet50"
    version = "V1"
    T_max = 10
    save_file = "./resultsDANet/" + current_time+"-"+str(is_balanced_weight) + ".txt"
    gpu_device = "0"
    cma_method = 'cpfnetsplit'           # cpfnet(883.24) cpfnetsplit(814.12) danet(1177.97) fuseda(795.47)
                                            # crossda(798.13) crossandselfda(851.29)
                                            # none(744.97) cbam(744.97) crosscbam(744.97) crossandselfcbam(744.97)
    fusion_method = 'addcatse'              # addcat addcatse centralnet cat

    this_time_log = f'''######################################################################
    Starting training:
    cma_method: {cma_method}
    fusion_method: {fusion_method}
    gpu: {gpu_device}
    network: {network}
    version: {version}
    save_file: {save_file}
    is_balanced_weight: {is_balanced_weight}
    img_size: {img_size}
    initial_learning_rate: {initial_learning_rate}
    momentum: {momentum}
    weight_decay: {weight_decay}
    train_batch_size: {train_batch_size}
    test_batch_size: {test_batch_size}
    num_workers: {num_workers}
######################################################################
    '''

    if is_balanced_weight:
        class_weight_diag = [1.0/42, 1.0/575, 1.0/252, 1.0/97, 1.0/45]
        class_weight_pn = [1.0/400, 1.0/381, 1.0/230]
        class_weight_bwv = [1.0/816, 1.0/195]
        class_weight_vs = [1.0/823, 1.0/117, 1.0/71]
        class_weight_pig = [1.0/588, 1.0/118, 1.0/305]
        class_weight_str = [1.0/653, 1.0/107, 1.0/251]
        class_weight_dag = [1.0/229, 1.0/334, 1.0/448]
        class_weight_rs = [1.0/758, 1.0/253]
    else:
        class_weight_diag = [1, 1, 1, 1, 1]
        class_weight_pn = [1, 1, 1]
        class_weight_bwv = [1, 1]
        class_weight_vs = [1, 1, 1]
        class_weight_pig = [1, 1, 1]
        class_weight_str = [1, 1, 1]
        class_weight_dag = [1, 1, 1]
        class_weight_rs = [1, 1]
    class_weight = [class_weight_pn, class_weight_str, class_weight_pig, class_weight_rs, class_weight_dag,
                    class_weight_bwv, class_weight_vs, class_weight_diag]

    basepath = r'/media/wz209/a29353b7-1090-433f-b452-b4ce827adb17/sugurs/Dataset/release_v0'
    data_train = os.path.join(basepath, r'meta/train_indexes.csv')
    data_val = os.path.join(basepath, r'meta/valid_indexes.csv')
    data_test = os.path.join(basepath, r'meta/test_indexes.csv')

    # data_train = '_train_indexes.txt'
    # data_val = '_valid_indexes.txt'
    # data_test = '_test_indexes.txt'

